package tagdb.algorithm.tag;

import com.compomics.util.Util;
import tagdb.fragmentation.FragmentIon;
import tagdb.fragmentation.Masses;
import tagdb.fragmentation.SpectrumPeak;
import tagdb.io.MascotGenericFile;
import tagdb.util.SpectrumUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * <b>TagScorer</b>
 * <p>
 * This class provides functions to score N- and C-terminal tags in
 * comparison experimental spectra.
 * </p>
 *
 * @author Thilo Muth
 */
public class TagScorer {

    /**
     * Tag to be used.
     */
    private Tag tag;

    /**
     * Peak list.
     */
    private List<SpectrumPeak> peakList;

    /**
     * Masses instance.
     */
    private Masses masses = Masses.getInstance();
    /**
     * The a ions.
     */
    private FragmentIon[] iAIons;
    /**
     * The a* ions.
     */
    private FragmentIon[] iANH3Ions;
    /**
     * The a° ions.
     */
    private FragmentIon[] iAH2OIons;
    /**
     * The b ions.
     */
    private FragmentIon[] iBIons;
    /**
     * The b* ions.
     */
    private FragmentIon[] iBNH3Ions;
    /**
     * The b° ions.
     */
    private FragmentIon[] iBH2OIons;
    /**
     * The c ions.
     */
    private FragmentIon[] iCIons;
    /**
     * The x ions.
     */
    private FragmentIon[] iXIons;
    /**
     * The y ions.
     */
    private FragmentIon[] iYIons;
    /**
     * The y* ions.
     */
    private FragmentIon[] iYNH3Ions;
    /**
     * The y° ions.
     */
    private FragmentIon[] iYH2OIons;
    /**
     * The z ions.
     */
    private FragmentIon[] iZIons;

    /**
     * List of theoretical ions.
     */
    private List<FragmentIon> theoIons;

    /**
     * Scored tag.
     */
    private double score;

    /**
     * Constructs the scoring algorithm for the tags.
     * @param tag Tag to be scored.
     * @param spectrum The experimental spectrum.
     */
    public TagScorer(Tag tag, MascotGenericFile spectrum) {
        this.tag = tag;
        this.peakList = SpectrumUtils.getPeakListFromSpectrum(spectrum);
        init();
        fragmentTag();
        matchIons(0.5);
    }

    /**
     * Initializes the tag scorer.
     */
    private void init() {
        int length = tag.getLength() * tag.getCharge();

        iAIons = new FragmentIon[length];
        iAH2OIons = new FragmentIon[length];
        iANH3Ions = new FragmentIon[length];
        iBIons = new FragmentIon[length];
        iBH2OIons = new FragmentIon[length];
        iBNH3Ions = new FragmentIon[length];
        iCIons = new FragmentIon[length];
        iXIons = new FragmentIon[length];
        iYIons = new FragmentIon[length];
        iYNH3Ions = new FragmentIon[length];
        iYH2OIons = new FragmentIon[length];
        iZIons = new FragmentIon[length];
    }

    /**
     * Fragments the peptide tag.
     */
    private void fragmentTag() {
        // Tag masses.
        double[] tagMasses = calculateTagMasses();
        double hydrogenMass = Masses.Hydrogen;
        double oxygenMass = Masses.Oxygen;
        double nitrogenMass = Masses.Nitrogen;
        double carbonMass = Masses.Carbon;
        double c_termMass = masses.get("C_term");

        int length = tag.getLength();

        int bIndex = 0;
        int yIndex = 0;
        for (int charge = 1; charge <= tag.getCharge(); charge++) {

            for (int i = 0; i < length; i++) {
                double bMass = 0.0;
                double yMass = 0.0;

                if(tag.getType() == Tag.TagType.N_TERMINAL) {
                    // Each peptide mass is added to the b ion mass
                    for (int j = 0; j <= i; j++) {
                        bMass += tagMasses[j];
                    }
                    // Create an instance for each fragment ion
                    if (charge <= tag.getCharge()) {
                        iAIons[bIndex] = new FragmentIon((bMass - oxygenMass - carbonMass + charge * hydrogenMass) / charge, FragmentIon.A_ION, i + 1, charge);
                        iANH3Ions[bIndex] = new FragmentIon((bMass - oxygenMass - carbonMass - nitrogenMass - 3 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.ANH3_ION, i + 1, charge);
                        iAH2OIons[bIndex] = new FragmentIon((bMass - 2 * oxygenMass - carbonMass - 2 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.AH2O_ION, i + 1, charge);
                        iBIons[bIndex] = new FragmentIon((bMass + charge * hydrogenMass) / charge, FragmentIon.B_ION, i + 1, charge);
                        iBNH3Ions[bIndex] = new FragmentIon((bMass - nitrogenMass - 3 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.BNH3_ION, i + 1, charge);
                        iBH2OIons[bIndex] = new FragmentIon((bMass - oxygenMass - 2 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.BH2O_ION, i + 1, charge);
                        iCIons[bIndex] = new FragmentIon((bMass + nitrogenMass + 3 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.C_ION, i + 1, charge);
                        bIndex++;
                    }
                }
                // C-terminal tag fragmentation.
                else if(tag.getType() == Tag.TagType.C_TERMINAL) {
                    // Each peptide mass is added to the y ion mass, taking the reverse direction (from the C terminal end)
                    for (int j = 0; j <= i; j++) {
                        yMass += tagMasses[(length - 1) - j];
                    }
                    // Add two extra hydrogen on the N terminal end and one hydroxyl at the C terminal end
                    yMass = yMass + c_termMass + hydrogenMass;

                    if (charge <= tag.getCharge()) {
                        // Create an instance of the fragment y ion
                        iXIons[yIndex] = new FragmentIon((yMass + carbonMass + oxygenMass - 2 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.X_ION, i + 1, charge);
                        iYIons[yIndex] = new FragmentIon((yMass + charge * hydrogenMass) / charge, FragmentIon.Y_ION, i + 1, charge);
                        iYNH3Ions[yIndex] = new FragmentIon((yMass - nitrogenMass - 3 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.YNH3_ION, i + 1, charge);
                        iYH2OIons[yIndex] = new FragmentIon((yMass - 2 * hydrogenMass - oxygenMass + charge * hydrogenMass) / charge, FragmentIon.YH2O_ION, i + 1, charge);
                        iZIons[yIndex] = new FragmentIon((yMass - nitrogenMass - 2 * hydrogenMass + charge * hydrogenMass) / charge, FragmentIon.Z_ION, i + 1, charge);
                        yIndex++;
                    }
                }
            }
        }
    }


    /**
     * This method calculates the masses of the tag, including the masses of the amino acids.
     *
     * @return double[] Contains the mass of the part of the tag sequence. The amino acid position is the index.
     */
    double[] calculateTagMasses() {
        double mass;

        String sequence = tag.getSequence();
        double[] peptideMasses = new double[tag.getLength()];
        for (int i = 0; i < sequence.length(); i++) {
            mass = 0.0;

            // For each amino acid add the specific mass
            String aa = String.valueOf(sequence.charAt(i));
            if (masses.containsKey(aa)) {
                mass += masses.get(aa);
                // Add each specific mass to the array
                peptideMasses[i] = mass;
            }
        }
        return peptideMasses;
    }

    /**
     * This method tries to match the theoretical masses of the ions with the
     * masses of the experimental peaks.
     *
     * @param fragMassTol The fragment mass tolerance
     * @return matchedIons A Vector containing all the matched fragment ions.
     */
    private List<FragmentIon> matchIons(double fragMassTol) {
        List<FragmentIon> matchedIons = new ArrayList<FragmentIon>();
        theoIons = new ArrayList<FragmentIon>();

        int start = 0 , end = 12;
        if(tag.getType() == Tag.TagType.N_TERMINAL) {
            start = 0;
            end = 6;
        }
        else if(tag.getType() == Tag.TagType.C_TERMINAL) {
            start = 7;
            end = 12;
        }

        for (int i = start; i < end; i++) {
            FragmentIon[] theoreticIons = getTheoreticIons(i);
            for (FragmentIon fragIon : theoreticIons) {
                if (fragIon != null) {
                    theoIons.add(fragIon);
                    if (fragIon.isMatch(peakList, fragMassTol)) {
                        matchedIons.add(fragIon);
                        score += fragIon.getScore();
                    }
                }
            }
        }
        tag.setScore(Util.roundDouble(score, 2));
        return matchedIons;
    }

    /**
     * Returns the corresponding array of theoretic ions.
     * 
     * @param type the ion type to return
     * @return the corresponding array of theoretic ions
     */
    public FragmentIon[] getTheoreticIons(int type) {
        switch (type) {
            case FragmentIon.A_ION:
                return iAIons;
            case FragmentIon.AH2O_ION:
                return iAH2OIons;
            case FragmentIon.ANH3_ION:
                return iANH3Ions;
            case FragmentIon.B_ION:
                return iBIons;
            case FragmentIon.BH2O_ION:
                return iBH2OIons;
            case FragmentIon.BNH3_ION:
                return iBNH3Ions;
            case FragmentIon.C_ION:
                return iCIons;
            case FragmentIon.X_ION:
                return iXIons;
            case FragmentIon.Y_ION:
                return iYIons;
            case FragmentIon.YH2O_ION:
                return iYH2OIons;
            case FragmentIon.YNH3_ION:
                return iYNH3Ions;
            case FragmentIon.Z_ION:
                return iZIons;
        }
        return null;
    }

    /**
     * Returns the calculated tag score.
     * @return Tag score.
     */
    public double getScore() {
        return score;
    }
}